import argparse
import h5py
import numpy as np
from utils import Image, glob_all, DicomImage
from fractal import multifractal_dimension
import argparse


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_folder', default='img/')
    parser.add_argument('--ext', choices=['*.jpg', '*.png', '*.IMA'])
    parser.add_argument('--bw_threshold', type=float, default=128.0)
    parser.add_argument('--bw_hu_threshold', type=float, default=0.0)
    parser.add_argument('--q_list', default="0,0.1,0.3,0.5")
    parser.add_argument('--block_height', type=int, default=40)
    parser.add_argument('--block_width', type=int, default=40)
    parser.add_argument('--height_step', type=int, default=40)
    parser.add_argument('--width_step', type=int, default=40)
    parser.add_argument('--output_file', default='data.csv')
    args = parser.parse_args()

    if args.ext == '*.IMA':
        args.is_dicom = True
    else:
        args.is_dicom = False

    args.q_list = [float(x) for x in args.q_list.strip().split(',')]

    return args


if __name__ == '__main__':
    args = parse_args()
    filenames = glob_all(args.input_folder, args.ext)
    with open(args.output_file, 'w') as f:
        for i, filename in enumerate(filenames):
            if args.is_dicom is True:
                img = DicomImage(filename, args.bw_hu_threshold)
            else:
                img = Image(filename, args.bw_threshold)
            for start_height, start_width, block_height, block_width, block in \
                    img.blockSetGenerator(args.block_height, args.block_width, args.height_step, args.width_step):
                features = multifractal_dimension(block, args.q_list)
                lines = [str(i), str(start_height), str(start_width), str(block_height),
                         str(block_width)] + [str(x) for x in features] + [filename]
                f.write(','.join(lines) + '\n')
            if (i + 1) % 1 == 0:
                print('{} / {} Finished.'.format(i + 1, len(filenames)))
